# 导入椭圆曲线算法
from ecdsa import SigningKey, SECP256k1, VerifyingKey, BadSignatureError
import binascii
import base64
from hashlib import sha256  # 用于哈希计算


class Wallet:
    """
    钱包类的定义
    """

    def __init__(self):
        """
        钱包初始化时基于椭圆曲线生成一个唯一的秘钥对，代表区块链上一个唯一的账户
        """
        self._private_key = SigningKey.generate(curve=SECP256k1)
        self._public_key = self._private_key.get_verifying_key()

    @property
    def address(self):  # TODO part 1
        """
        通过公钥生成地址（对公钥做 sha256 后 base64 编码）
        """
        pubkey_bytes = self._public_key.to_string()
        sha256_hash = sha256(pubkey_bytes).digest()
        return base64.b64encode(sha256_hash).decode('utf-8')

    @property
    def publicKey(self):
        """
        返回公钥字符串
        """
        return self._public_key.to_pem()

    def signature(self, message):  # TODO part 2
        """
        生成消息的数字签名
        """
        message_hash = sha256(message.encode('utf-8')).digest()
        signature = self._private_key.sign(message_hash)
        return base64.b64encode(signature).decode('utf-8')


def verify_signature(publicKey, message, signature):  # TODO part 3
    """
    验证签名
    """
    try:
        # 解析 PEM 格式的公钥
        vk = VerifyingKey.from_pem(publicKey)
        message_hash = sha256(message.encode('utf-8')).digest()
        signature_bytes = base64.b64decode(signature)
        return vk.verify(signature_bytes, message_hash)
    except (BadSignatureError, Exception):
        return False
